
% 初始化参数
clear all; close all; clc;

% 定义数据路径
imagesDir = 'images';
labelsDir = 'labels';
trainImagesDir = fullfile(imagesDir, 'train');
valImagesDir = fullfile(imagesDir, 'val');
trainLabelsDir = fullfile(labelsDir, 'train');
valLabelsDir = fullfile(labelsDir, 'val');

% 标签
classNames = {'axle', 'scratch', 'pit', 'bruise'};
numClasses = length(classNames);
% 图像大小
inputSize = [224, 224, 3]; 
% 准备数据
disp('正在加载训练和验证数据...');

% 训练数据
trainImages = imageDatastore(trainImagesDir, 'IncludeSubfolders', true, 'LabelSource', 'foldernames');
trainImageCount = numel(trainImages.Files);

% 验证数据
valImages = imageDatastore(valImagesDir, 'IncludeSubfolders', true, 'LabelSource', 'foldernames');
valImageCount = numel(valImages.Files);

% 加载YOLO标签
trainLabels = cell(trainImageCount, 1);
valLabels = cell(valImageCount, 1);

% 加载训练标签
for i = 1:trainImageCount
    [~, imageName, ~] = fileparts(trainImages.Files{i});
    labelFile = fullfile(trainLabelsDir, [imageName, '.txt']);
    if exist(labelFile, 'file')
        trainLabels{i} = loadYoloLabels(labelFile);
    else
        warning('标签文件未找到: %s', labelFile);
        trainLabels{i} = [];
    end
end

% 加载验证标签
for i = 1:valImageCount
    [~, imageName, ~] = fileparts(valImages.Files{i});
    labelFile = fullfile(valLabelsDir, [imageName, '.txt']);
    if exist(labelFile, 'file')
        valLabels{i} = loadYoloLabels(labelFile);
    else
        warning('标签文件未找到: %s', labelFile);
        valLabels{i} = [];
    end
end

% 特征提取和融合
disp('正在提取多模态特征...');
trainFeatures = {};
validImagePaths = {};
trainClassLabels = [];

% 数据初始化
valImagePaths = {};
valClassLabels = [];

% 为每个训练图像提取特征
for i = 1:trainImageCount
    try
        img = readimage(trainImages, i);
        img = imresize(img, inputSize(1:2));
        
        if ~isempty(trainLabels{i})
            % 获取标签信息
            bbox = trainLabels{i}(2:5);
            classId = trainLabels{i}(1) + 1; 
            [x, y, width, height] = yoloToPixel(bbox, size(img));
            roi = imcrop(img, [x, y, width, height]);
            if ~isempty(roi)
                % 提取多模态特征
                features = extractFeatures(roi);
                
                % 保存特征
                trainFeatures{end+1} = features;
                validImagePaths{end+1} = trainImages.Files{i};
                trainClassLabels(end+1) = classId;
            end
        end
    catch e
        warning('处理图像 %d 时出错: %s', i, e.message);
    end
end

% 验证数据处理代码
disp('正在处理验证数据...');

% 处理验证图像
for i = 1:valImageCount
    try
        img = readimage(valImages, i);
        img = imresize(img, inputSize(1:2));
        
        if ~isempty(valLabels{i})
            % 获取标签信息
            bbox = valLabels{i}(2:5);
            classId = valLabels{i}(1) + 1;
            [x, y, width, height] = yoloToPixel(bbox, size(img));
            roi = imcrop(img, [x, y, width, height]);
            if ~isempty(roi)
                % 保存相关信息
                valImagePaths{end+1} = valImages.Files{i};
                valClassLabels(end+1) = classId;
            end
        end
    catch e
        warning('处理验证图像 %d 时出错: %s', i, e.message);
    end
end

% 定义数据增强器
disp('处理数据增强...');
% 创建数据增强选项
augmenter = imageDataAugmenter( ...
    'RandRotation', [-30, 30], ...
    'RandXScale', [0.8, 1.2], ...
    'RandYScale', [0.8, 1.2], ...
    'RandXShear', [-0.2, 0.2], ...
    'RandYShear', [-0.2, 0.2]);

%%创建训练数据
tempTrainDir = fullfile(pwd, 'temp_train_images');
if ~exist(tempTrainDir, 'dir')
    mkdir(tempTrainDir);
end
for c = 1:numClasses
    classDir = fullfile(tempTrainDir, classNames{c});
    if ~exist(classDir, 'dir')
        mkdir(classDir);
    end
end
for i = 1:length(validImagePaths)
    [~, fileName, fileExt] = fileparts(validImagePaths{i});
    sourceFile = validImagePaths{i};
    classIdx = trainClassLabels(i);
    targetDir = fullfile(tempTrainDir, classNames{classIdx});
    targetFile = fullfile(targetDir, [fileName, fileExt]);
    copyfile(sourceFile, targetFile);
end

% 数据存储
trainDS = imageDatastore(tempTrainDir, 'IncludeSubfolders', true, 'LabelSource', 'foldernames');
augimdsTrain = augmentedImageDatastore(inputSize(1:2), trainDS, 'DataAugmentation', augmenter);

%验证图像
tempValDir = fullfile(pwd, 'temp_val_images');
if ~exist(tempValDir, 'dir')
    mkdir(tempValDir);
end
for c = 1:numClasses
    classDir = fullfile(tempValDir, classNames{c});
    if ~exist(classDir, 'dir')
        mkdir(classDir);
    end
end
for i = 1:length(valImagePaths)
    [~, fileName, fileExt] = fileparts(valImagePaths{i});
    sourceFile = valImagePaths{i};
    classIdx = valClassLabels(i);
    targetDir = fullfile(tempValDir, classNames{classIdx});
    targetFile = fullfile(targetDir, [fileName, fileExt]);
    copyfile(sourceFile, targetFile);
end

% 样本数量
for c = 1:numClasses
    classDir = fullfile(tempValDir, classNames{c});
    files = dir(fullfile(classDir, '*.jpg'));
    
    if length(files) == 0
        sourceDir = fullfile(tempTrainDir, classNames{c});
        sourceFiles = dir(fullfile(sourceDir, '*.jpg'));
        if ~isempty(sourceFiles)
            numToCopy = min(5, length(sourceFiles));
            for j = 1:numToCopy
                sourceFile = fullfile(sourceDir, sourceFiles(j).name);
                targetFile = fullfile(classDir, sourceFiles(j).name);
                copyfile(sourceFile, targetFile);
            end
        end
    end
end

% 创建验证数据存储
valDS = imageDatastore(tempValDir, 'IncludeSubfolders', true, 'LabelSource', 'foldernames');
augimdsVal = augmentedImageDatastore(inputSize(1:2), valDS);

% 设置迁移学习CNN
disp('正在设置迁移学习CNN模型...');
% 加载预训练的ResNet50
net = resnet50();
lgraph = layerGraph(net);
layersToRemove = {
    'fc1000'
    'fc1000_softmax'
    'ClassificationLayer_fc1000'
};

lgraph = removeLayers(lgraph, layersToRemove);

% 分类层
newLayers = [
    fullyConnectedLayer(numClasses, 'Name', 'fc_defects', 'WeightLearnRateFactor', 10, 'BiasLearnRateFactor', 10)
    softmaxLayer('Name', 'softmax')
    classificationLayer('Name', 'classoutput')
];

lgraph = addLayers(lgraph, newLayers);

% 连接新层
lgraph = connectLayers(lgraph, 'avg_pool', 'fc_defects');
% 训练CNN
disp('开始训练CNN模型...');

% 设置训练选项
options = trainingOptions('sgdm', ...
    'MiniBatchSize', 32, ...
    'MaxEpochs', 3, ...
    'InitialLearnRate', 0.001, ...
    'LearnRateSchedule', 'piecewise', ...
    'LearnRateDropFactor', 0.1, ...
    'LearnRateDropPeriod', 10, ...
    'Shuffle', 'every-epoch', ...
    'ValidationData', augimdsVal, ...
    'ValidationFrequency', 50, ...
    'Verbose', false, ...
    'Plots', 'training-progress', ...
    'ExecutionEnvironment', 'auto');

% 训练网络
net = trainNetwork(augimdsTrain, lgraph, options);

% 评估模型
disp('评估模型性能...');
[YPred, scores] = classify(net, augimdsVal);
valTrueLabels = valDS.Labels;
targetClass = 'axle';
targetClassIndex = find(ismember(classNames, targetClass));

% 提取样本
axleTrueLabels = (valTrueLabels == categorical({'axle'}));
axleScores = scores(:, targetClassIndex);

% 计算性能指标
TP = sum(YPred == categorical({'axle'}) & valTrueLabels == categorical({'axle'}));
FP = sum(YPred == categorical({'axle'}) & valTrueLabels ~= categorical({'axle'}));
FN = sum(YPred ~= categorical({'axle'}) & valTrueLabels == categorical({'axle'}));
precision = TP / (TP + FP);
recall = TP / (TP + FN);
f1 = 2 * (precision * recall) / (precision + recall);

% 输出
disp(['评估性能指标:']);
disp(['  精确率 = ', num2str(precision * 100), '%']);
disp(['  召回率 = ', num2str(recall * 100), '%']);
disp(['  F1分数 = ', num2str(f1)]);

% 绘制ROC 曲线
figure('Name', 'ROC曲线', 'Position', [100, 600, 600, 500]);
[X, Y, T, AUC] = perfcurve(axleTrueLabels, axleScores, true);
plot(X, Y, 'LineWidth', 2);
hold on;
plot([0, 1], [0, 1], 'r--');
hold off;

xlabel('假阳性率');
ylabel('真阳性率');
title(['ROC曲线']);
grid on;
hold off;
disp('缺陷检测系统训练完成!'); 